{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/zhangzhengming/下载/yolov3/data\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'data/ImageSets/train.txt'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-1-fee732128917>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m()\u001b[0m\n\u001b[1;32m     51\u001b[0m     \u001b[0;32mif\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mpath\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mexists\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data/labels/'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     52\u001b[0m         \u001b[0mos\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mmakedirs\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data/labels/'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 53\u001b[0;31m     \u001b[0mimage_ids\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data/ImageSets/%s.txt'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mimage_set\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mread\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstrip\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msplit\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     54\u001b[0m     \u001b[0mlist_file\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mopen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'data/%s.txt'\u001b[0m \u001b[0;34m%\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mimage_set\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m'w'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     55\u001b[0m     \u001b[0;32mfor\u001b[0m \u001b[0mimage_id\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mimage_ids\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'data/ImageSets/train.txt'"
     ]
    }
   ],
   "source": [
    "import xml.etree.ElementTree as ET\n",
    "import pickle\n",
    "import os\n",
    "from os import listdir, getcwd\n",
    "from os.path import join\n",
    " \n",
    "sets = ['train', 'test','val']\n",
    " \n",
    "classes = [\"collector\"]\n",
    " \n",
    " \n",
    "def convert(size, box):\n",
    "    dw = 1. / size[0]\n",
    "    dh = 1. / size[1]\n",
    "    x = (box[0] + box[1]) / 2.0\n",
    "    y = (box[2] + box[3]) / 2.0\n",
    "    w = box[1] - box[0]\n",
    "    h = box[3] - box[2]\n",
    "    x = x * dw\n",
    "    w = w * dw\n",
    "    y = y * dh\n",
    "    h = h * dh\n",
    "    return (x, y, w, h)\n",
    " \n",
    " \n",
    "def convert_annotation(image_id):\n",
    "    in_file = open('Annotations/%s.xml' % (image_id))\n",
    "    out_file = open('data/labels/%s.txt' % (image_id), 'w')\n",
    "    tree = ET.parse(in_file)\n",
    "    root = tree.getroot()\n",
    "    size = root.find('size')\n",
    "    w = int(size.find('width').text)\n",
    "    h = int(size.find('height').text)\n",
    " \n",
    "    for obj in root.iter('object'):\n",
    "        difficult = obj.find('Difficult').text\n",
    "        cls = obj.find('name').text\n",
    "        if cls not in classes or int(difficult) == 1:\n",
    "            continue\n",
    "        cls_id = classes.index(cls)\n",
    "        xmlbox = obj.find('bndbox')\n",
    "        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),\n",
    "             float(xmlbox.find('ymax').text))\n",
    "        bb = convert((w, h), b)\n",
    "        out_file.write(str(cls_id) + \" \" + \" \".join([str(a) for a in bb]) + '\\n')\n",
    " \n",
    " \n",
    "wd = getcwd()\n",
    "print(wd)\n",
    "for image_set in sets:\n",
    "    if not os.path.exists('data/labels/'):\n",
    "        os.makedirs('data/labels/')\n",
    "    image_ids = open('data/ImageSets/%s.txt' % (image_set)).read().strip().split()\n",
    "    list_file = open('data/%s.txt' % (image_set), 'w')\n",
    "    for image_id in image_ids:\n",
    "        list_file.write('images/%s.jpg\\n' % (image_id))\n",
    "        convert_annotation(image_id)\n",
    "    list_file.close()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
